import os
import subprocess
import sys
from typing import Optional
from datetime import datetime
import paramiko
import yaml



class Deployer:
    """
    部署工具类，用于构建Docker镜像、推送到阿里云镜像仓库并在远端服务器启动容器
    """

    def __init__(self,
                 image_name: str,
                 image_tag: str = "latest",
                 registry_url: Optional[str] = None,
                 remote_host: Optional[str] = None,
                 remote_password: Optional[str] = None,
                 remote_user: Optional[str] = None):
        """
        初始化部署器

        Args:
            image_name: 镜像名称
            image_tag: 镜像标签
            registry_url: 镜像仓库地址
            remote_host: 远端服务器地址
            remote_user: 远端服务器用户名
        """
        self.image_name = image_name
        self.image_tag = image_tag
        self.registry_url = registry_url
        self.remote_host = remote_host
        self.remote_password = remote_password
        self.remote_user = remote_user
        self.full_image_name = f"{registry_url}/{image_name}:{image_tag}" if registry_url else f"{image_name}:{image_tag}"

    def build_image(self, dockerfile_path: str = "Dockerfile", context_path: str = ".", build_dir: str = None) -> bool:
        """
        在本地根据Dockerfile构建镜像

        Args:
            dockerfile_path: Dockerfile文件路径（相对于context_path）
            context_path: 构建上下文路径
            build_dir: 构建执行目录（如果指定，则在此目录下执行构建命令）

        Returns:
            bool: 构建是否成功
        """
        print(f"开始构建镜像: {self.full_image_name}")

        # 确定实际的工作目录
        working_dir = build_dir if build_dir else "."
        print(f"在目录 {working_dir} 中执行构建")
        print(f"构建上下文路径: {context_path}")

        try:
            command = [
                "docker", "build",
                "-t", self.full_image_name,
                "-f", dockerfile_path,
                context_path
            ]

            print(f"执行命令: {' '.join(command)}")

            # 如果指定了build_dir，则在该目录下执行命令
            result = subprocess.run(command, capture_output=True, text=True,
                                    encoding="utf-8",
                                    cwd=working_dir if build_dir else None)

            if result.returncode == 0:
                print("镜像构建成功")
                if result.stdout:
                    print(result.stdout)
                return True
            else:
                print("镜像构建失败:")
                print(result.stderr)
                return False
        except Exception as e:
            print(f"构建镜像时发生异常: {e}")
            return False

    def push_image(self) -> bool:
        """
        将镜像推送到阿里云镜像仓库

        Returns:
            bool: 推送是否成功
        """
        print(f"开始推送镜像到阿里云镜像仓库: {self.full_image_name}")
        try:
            command = ["docker", "push", self.full_image_name]
            print(f"执行命令: {' '.join(command)}")

            result = subprocess.run(command, capture_output=True, text=True)

            if result.returncode == 0:
                print("镜像推送成功")
                if result.stdout:
                    print(result.stdout)
                return True
            else:
                print("镜像推送失败:")
                print(result.stderr)
                return False
        except Exception as e:
            print(f"推送镜像时发生异常: {e}")
            return False

    def update_compose_file(self, compose_file_path: str = "docker-compose.yml",
                            service_name: str = "backend") -> str:
        """
        更新docker-compose.yml文件中的镜像版本号

        Args:
            compose_file_path: docker-compose.yml文件路径
            service_name: 服务名称

        Returns:
            str: 更新后的临时文件路径
        """
        print(f"更新 {compose_file_path} 中的镜像版本号为 {self.full_image_name}")

        try:
            # 读取原始docker-compose文件
            with open(compose_file_path, 'r', encoding='utf-8') as file:
                compose_data = yaml.safe_load(file)

            # 更新指定服务的镜像版本
            if 'services' in compose_data and service_name in compose_data['services']:
                compose_data['services'][service_name]['image'] = self.full_image_name
                print(f"已更新服务 {service_name} 的镜像为: {self.full_image_name}")
            else:
                print(f"警告: 在 {compose_file_path} 中未找到服务 {service_name}")

            # 创建临时文件
            temp_file_path = f"temp_docker-compose_{self.image_tag}.yml"

            # 写入更新后的配置
            with open(temp_file_path, 'w', encoding='utf-8') as file:
                yaml.dump(compose_data, file, default_flow_style=False, allow_unicode=True,
                          indent=2, sort_keys=False)

            print(f"已生成更新后的compose文件: {temp_file_path}")
            return temp_file_path

        except Exception as e:
            print(f"更新docker-compose文件时发生异常: {e}")
            return compose_file_path  # 返回原始文件路径

    def _execute_remote_commands_paramiko(self, commands: list, compose_project_dir: str) -> tuple[bool, str, str]:
        """
        使用 paramiko 在远端服务器执行命令

        Args:
            commands: 要执行的命令列表
            compose_project_dir: docker-compose.yml文件所在目录

        Returns:
            tuple: (是否成功, stdout, stderr)
        """
        try:
            # 创建SSH客户端
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            # 使用密码连接
            ssh.connect(self.remote_host, username=self.remote_user,
                        password=self.remote_password)

            # 切换到工作目录并执行命令
            full_command = f"cd {compose_project_dir} && " + " && ".join(commands)
            print(f"将在远端执行命令: {full_command}")

            # 执行命令
            stdin, stdout, stderr = ssh.exec_command(full_command)

            # 获取执行结果
            output = stdout.read().decode('utf-8')
            error = stderr.read().decode('utf-8')
            exit_status = stdout.channel.recv_exit_status()

            ssh.close()

            return exit_status == 0, output, error

        except Exception as e:
            return False, "", str(e)

    def deploy_to_remote(self, service_name: str = "backend",
                         compose_file_path: str = "docker-compose.yml",
                         compose_project_dir: str = "/opt/code",
                         remote_compose: str = "docker-compose.yml") -> bool:
        """
        在远端服务器使用Docker Compose部署并启动特定服务

        Args:
            service_name: docker-compose.yml中定义的服务名称
            compose_file_path: 本地docker-compose.yml文件路径
            compose_project_dir: 远端服务器docker-compose.yml文件所在目录
            remote_compose:
        Returns:
            bool: 部署是否成功

        """
        print(f"开始在远端服务器 {self.remote_user}@{self.remote_host} 部署应用")
        print(f"使用Docker Compose在目录 {compose_project_dir} 中部署服务 {service_name}")

        # 1. 更新docker-compose文件中的镜像版本
        updated_compose_file = self.update_compose_file(compose_file_path=compose_file_path, service_name=service_name)

        # 2. 上传更新后的docker-compose文件到服务器
        remote_compose_path = f"{compose_project_dir}/{remote_compose}"
        if not self._upload_file_paramiko(updated_compose_file, remote_compose_path):
            print("上传docker-compose文件失败")
            # 清理临时文件
            if updated_compose_file != compose_file_path and os.path.exists(updated_compose_file):
                os.remove(updated_compose_file)
            return False

        # 3. 清理临时文件
        if updated_compose_file != compose_file_path and os.path.exists(updated_compose_file):
            os.remove(updated_compose_file)

        # 4. 在服务器上执行部署命令
        remote_commands = [
            f"docker pull {self.full_image_name}",
            f"docker compose -f {remote_compose} down",
            f"docker compose -f {remote_compose} up -d"
        ]

        # 使用 paramiko 执行远程命令
        success, stdout, stderr = self._execute_remote_commands_paramiko(remote_commands, compose_project_dir)

        if success:
            print("远端部署成功")
            if stdout:
                print(stdout)
            return True
        else:
            print("远端部署失败:")
            if stderr:
                print(stderr)
            return False

    def _upload_file_paramiko(self, local_path: str, remote_path: str) -> bool:
        """
        使用 paramiko 上传文件到远端服务器

        Args:
            local_path: 本地文件路径
            remote_path: 远端文件路径

        Returns:
            bool: 上传是否成功
        """
        try:
            # 创建SSH客户端
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            # 使用密码连接
            ssh.connect(self.remote_host, username=self.remote_user,
                        password=self.remote_password)

            # 创建SFTP客户端
            sftp = ssh.open_sftp()

            # 上传文件
            print(f"上传文件: {local_path} -> {remote_path}")
            sftp.put(local_path, remote_path)

            # 关闭连接
            sftp.close()
            ssh.close()

            print("文件上传成功")
            return True

        except Exception as e:
            print(f"文件上传失败: {e}")
            return False

    def tag_image(self, source_image: str) -> bool:
        """
        为镜像打标签

        Args:
            source_image: 源镜像名称

        Returns:
            bool: 打标签是否成功
        """
        print(f"为镜像 {source_image} 打标签: {self.full_image_name}")
        try:
            command = ["docker", "tag", source_image, self.full_image_name]
            result = subprocess.run(command, capture_output=True, text=True)

            if result.returncode == 0:
                print("镜像标签添加成功")
                return True
            else:
                print("镜像标签添加失败:")
                print(result.stderr)
                return False
        except Exception as e:
            print(f"为镜像打标签时发生异常: {e}")
            return False

    def cleanup_old_images_and_containers(self) -> bool:
        """
        清理本地多余的镜像和停止的容器

        Returns:
            bool: 清理是否成功
        """
        print("开始清理本地多余的镜像和容器...")
        try:
            # 清理停止的容器
            print("清理停止的容器...")
            prune_containers_cmd = ["docker", "container", "prune", "-f"]
            result = subprocess.run(prune_containers_cmd, capture_output=True, text=True)
            if result.returncode == 0:
                print("停止的容器清理完成")
            else:
                print(f"清理停止的容器时出错: {result.stderr}")

            # 清理未使用的镜像（不包括当前标签的镜像）
            print("清理未使用的镜像...")
            prune_images_cmd = ["docker", "image", "prune", "-a", "-f"]
            result = subprocess.run(prune_images_cmd, capture_output=True, text=True)
            if result.returncode == 0:
                print("未使用的镜像清理完成")
            else:
                print(f"清理未使用的镜像时出错: {result.stderr}")

            # 可选：列出当前相关的镜像
            print("当前相关的镜像:")
            list_images_cmd = ["docker", "images", self.image_name]
            subprocess.run(list_images_cmd, capture_output=True, text=True)

            print("本地镜像和容器清理完成!")
            return True
        except Exception as e:
            print(f"清理过程中发生异常: {e}")
            return False

    def cleanup_remote_old_images_and_containers(self) -> bool:
        """
        清理远程服务器上多余的镜像和停止的容器

        Returns:
            bool: 清理是否成功
        """
        print("开始清理远程服务器上多余的镜像和容器...")
        try:
            # 在远程服务器上执行清理命令
            remote_commands = [
                "docker container prune -f",  # 清理停止的容器
                "docker image prune -a -f"    # 清理未使用的镜像
            ]
            
            # 使用 paramiko 执行远程命令
            success, stdout, stderr = self._execute_remote_commands_paramiko(remote_commands, "/")
            
            if success:
                print("远程服务器镜像和容器清理完成!")
                if stdout:
                    print(stdout)
                return True
            else:
                print("远程服务器清理过程中出错:")
                if stderr:
                    print(stderr)
                return False
                
        except Exception as e:
            print(f"远程清理过程中发生异常: {e}")
            return False


# 配置参数 - 根据你的实际环境修改
IMAGE_NAME = "priv-backend"
IMAGE_NAME_FRT = "priv-frontend"
IMAGE_TAG = f"v{datetime.now().strftime('%Y%m%d%H%M%S')}"  # 使用时间戳作为版本号
# IMAGE_TAG = "v1.1"
REGISTRY_URL = "crpi-lt6kndkur5fb93d0.cn-hangzhou.personal.cr.aliyuncs.com/qnn-25"
REMOTE_HOST = "8.155.160.155"  # 修改为你的服务器地址
REMOTE_USER = "root"  # 修改为你的服务器用户


def deploy_backend():
    """
    部署后端代码
    :return:
    """
    run_deploy(IMAGE_NAME)


def deploy_frontend():
    build_dir = "D:\\code\\priv_domain_web"
    compose_file_path = "D:\\code\\priv_domain_web\\docker-compose-ft.yml"
    remote_compose = "docker-compose-ft.yml"
    kwargs = {
        "compose_file_path": compose_file_path,
        "remote_compose": remote_compose,
        "service_name": "frontend"
    }
    run_deploy(IMAGE_NAME_FRT, build_dir=build_dir, **kwargs)


def run_deploy(image_name, build_dir=None, **kwargs):
    print(f"开始部署 {image_name} 应用，版本: {IMAGE_TAG}")

    # 创建部署器实例
    deployer = Deployer(image_name=image_name, image_tag=IMAGE_TAG, registry_url=REGISTRY_URL,
                        remote_host=REMOTE_HOST, remote_user=REMOTE_USER, remote_password="Zz!111111")

    # 构建镜像
    if not deployer.build_image(build_dir=build_dir):
        print("构建镜像失败，终止部署")
        sys.exit(1)

    # 推送镜像
    if not deployer.push_image():
        print("推送镜像失败，终止部署")
        sys.exit(1)

    # 部署到远端服务器
    success = deployer.deploy_to_remote(**kwargs)

    if success:
        print("部署完成!")
        print(f"应用已部署为版本: {IMAGE_TAG}")
        
        # 清理本地多余的镜像和容器
        print("开始清理本地多余的镜像和容器...")
        cleanup_success = deployer.cleanup_old_images_and_containers()
        if cleanup_success:
            print("本地镜像和容器清理完成!")
        else:
            print("清理过程中出现警告，但不影响部署结果")
            
        # 清理远程服务器上多余的镜像和容器
        print("开始清理远程服务器上多余的镜像和容器...")
        remote_cleanup_success = deployer.cleanup_remote_old_images_and_containers()
        if remote_cleanup_success:
            print("远程服务器镜像和容器清理完成!")
        else:
            print("远程服务器清理过程中出现警告，但不影响部署结果")
    else:
        print("部署失败!")
        sys.exit(1)


if __name__ == "__main__":
    deploy_frontend()
    # deploy_backend()